package com.reebake.ideal.protect.filter;

import com.reebake.ideal.crypto.core.SecretKeyService;
import com.reebake.ideal.crypto.entity.SecretKeyEntity;
import com.reebake.ideal.protect.core.AbstractProtectFilter;
import com.reebake.ideal.protect.core.ContentSecureService;
import com.reebake.ideal.protect.core.ProtectParameter;
import com.reebake.ideal.protect.properties.ProtectProperties;
import com.reebake.ideal.protect.web.DecryptedRequestWrapper;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.http.HttpHeaders;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.util.ContentCachingResponseWrapper;

import java.io.IOException;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

/**
 * 内容保护过滤器，请求解密和输出加密
 */
public class ContentProtectFilter extends AbstractProtectFilter {
    private final ContentSecureService contentSecureService;
    private final SecretKeyService secretKeyService;

    public ContentProtectFilter(ProtectProperties protectProperties, ContentSecureService contentSecureService, SecretKeyService secretKeyService) {
        super(protectProperties);
        this.secretKeyService = secretKeyService;
        this.contentSecureService = contentSecureService;
    }

    @Override
    protected void protect(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
        ProtectParameter protectParameter = getProtectParameter(request);
        SecretKeyEntity secretKeyEntity = secretKeyService.queryByAppKey(protectParameter.getAppKey());
        DecryptedRequestWrapper requestWrapper = new DecryptedRequestWrapper(request, contentSecureService, secretKeyEntity);
        ContentCachingResponseWrapper responseWrapper = new ContentCachingResponseWrapper(response);
        ServletRequestAttributes requestAttributes = new ServletRequestAttributes(requestWrapper);
        // 替换原来的请求，使多次读取内容一致
        RequestContextHolder.setRequestAttributes(requestAttributes);

        filterChain.doFilter(requestWrapper, responseWrapper);

        // 重置响应并加密响应内容
        byte[] content = responseWrapper.getContentAsByteArray();
        String originalResponseBody = new String(content);
        String encryptedResponseBody = contentSecureService.encrypt(originalResponseBody, secretKeyEntity);
        // 写响应头
        Map<String, String> headers = new HashMap<>();
        int status = responseWrapper.getStatus();
        Enumeration<String> enumeration = requestWrapper.getHeaderNames();
        while (enumeration.hasMoreElements()) {
            String name = enumeration.nextElement();
            if(HttpHeaders.CONTENT_LENGTH.equalsIgnoreCase(name)) {
                continue;
            }
            headers.put(name, responseWrapper.getHeader(name));
        }

        // 写回客户端
        responseWrapper.reset();
        responseWrapper.setStatus(status);
        for(Map.Entry<String, String> entry : headers.entrySet()) {
            responseWrapper.setHeader(entry.getKey(), entry.getValue());
        }
        responseWrapper.getWriter().write(encryptedResponseBody);
        responseWrapper.copyBodyToResponse();

    }
}
